#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
验证 AIGC 元数据
"""

import json
import sys
import os

# 设置输出编码为 UTF-8
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

try:
    from PIL import Image
    from PIL.ExifTags import TAGS
except ImportError:
    print("错误: 需要安装 Pillow")
    sys.exit(1)

try:
    from pyexiv2 import Image as ExivImage
    HAS_PYEXIV2 = True
except ImportError:
    HAS_PYEXIV2 = False


def read_png_metadata(image_path):
    """
    读取 PNG tEXt chunk 元数据
    """
    try:
        img = Image.open(image_path)
        
        # PNG 使用 text 字典存储元数据
        if hasattr(img, 'text'):
            # 尝试读取 Comment
            if 'Comment' in img.text:
                comment = img.text['Comment']
                if comment.startswith('aigc:'):
                    json_str = comment[5:]  # 去掉 'aigc:' 前缀
                    return json.loads(json_str)
            
            # 尝试读取 AIGC
            if 'AIGC' in img.text:
                return json.loads(img.text['AIGC'])
        
        # 尝试从 info 中读取
        if hasattr(img, 'info'):
            if 'Comment' in img.info:
                comment = img.info['Comment']
                if isinstance(comment, bytes):
                    comment = comment.decode('utf-8', errors='ignore')
                if comment.startswith('aigc:'):
                    json_str = comment[5:]
                    return json.loads(json_str)
            
            if 'AIGC' in img.info:
                aigc = img.info['AIGC']
                if isinstance(aigc, bytes):
                    aigc = aigc.decode('utf-8', errors='ignore')
                return json.loads(aigc)
        
        return None
    except Exception as e:
        print(f"[WARN] 读取 PNG 元数据失败: {e}")
        return None


def read_xmp_metadata(image_path):
    """
    读取 XMP 元数据
    """
    if not HAS_PYEXIV2:
        return None
    
    try:
        img = ExivImage(image_path)
        
        # 读取 XMP.dc.aigc
        xmp_data = img.read_xmp()
        
        # 查找 aigc 相关字段
        aigc_data = {}
        
        # 尝试读取 Xmp.dc.aigc
        if 'Xmp.dc.aigc' in xmp_data:
            try:
                return json.loads(xmp_data['Xmp.dc.aigc'])
            except:
                pass
        
        # 尝试从 description 中提取
        if 'Xmp.dc.description' in xmp_data:
            desc = xmp_data['Xmp.dc.description']
            if desc.startswith('AIGC: '):
                try:
                    return json.loads(desc[7:])
                except:
                    pass
        
        img.close()
        return None
        
    except Exception as e:
        print(f"[WARN] 读取 XMP 失败: {e}")
        return None


def main():
    image_path = os.path.expanduser("~/Documents/test/testPng.png")
    
    if not os.path.exists(image_path):
        print(f"错误: 图片不存在: {image_path}")
        sys.exit(1)
    
    print("=" * 60)
    print("验证 AIGC 元数据 (PNG)")
    print("=" * 60)
    print(f"\n图片路径: {image_path}\n")
    
    # 读取 PNG tEXt 元数据
    print("1. 读取 PNG tEXt 元数据:")
    png_data = read_png_metadata(image_path)
    if png_data:
        print(json.dumps(png_data, indent=2, ensure_ascii=False))
    else:
        print("   [未找到]")
    
    print("\n2. 读取 XMP 元数据:")
    xmp_data = read_xmp_metadata(image_path)
    if xmp_data:
        print(json.dumps(xmp_data, indent=2, ensure_ascii=False))
    else:
        print("   [未找到]")
    
    print("\n" + "=" * 60)
    print("验证完成")
    print("=" * 60)


if __name__ == "__main__":
    main()

